Tensor-Vector Outer Product using Einsum

This is an example of how to properly use the einsum function to compute a tensor-vector outer product.

from csdl_om import Simulatorimport numpy as npfrom csdl import Modelimport csdl

class ExampleOuterTensorVector(Model):
    def define(self):        a = np.arange(4)        vec = self.declare_variable('a', val=a)
        # Shape of Tensor        shape3 = (2, 4, 3)        c = np.arange(24).reshape(shape3)
        # Declaring tensor        tens = self.declare_variable('c', val=c)
        # Outer Product of a tensor and a vector        self.register_output(            'einsum_outer2',            csdl.einsum_new_api(                tens,                vec,                operation=[(0, 1, 30), (2, ), (0, 1, 30, 2)],            ))

sim = Simulator(ExampleOuterTensorVector())sim.run()
print('a', sim['a'].shape)print(sim['a'])print('c', sim['c'].shape)print(sim['c'])print('einsum_outer2', sim['einsum_outer2'].shape)print(sim['einsum_outer2'])

[0. 1. 2. 3.]c (2, 4, 3)[[[ 0.  1.  2.]  [ 3.  4.  5.]  [ 6.  7.  8.]  [ 9. 10. 11.]]
 [[12. 13. 14.]  [15. 16. 17.]  [18. 19. 20.]  [21. 22. 23.]]]einsum_outer2 (2, 4, 3, 4)[[[[ 0.  0.  0.  0.]   [ 0.  1.  2.  3.]   [ 0.  2.  4.  6.]]
  [[ 0.  3.  6.  9.]   [ 0.  4.  8. 12.]   [ 0.  5. 10. 15.]]
  [[ 0.  6. 12. 18.]   [ 0.  7. 14. 21.]   [ 0.  8. 16. 24.]]
  [[ 0.  9. 18. 27.]   [ 0. 10. 20. 30.]   [ 0. 11. 22. 33.]]]

 [[[ 0. 12. 24. 36.]   [ 0. 13. 26. 39.]   [ 0. 14. 28. 42.]]
  [[ 0. 15. 30. 45.]   [ 0. 16. 32. 48.]   [ 0. 17. 34. 51.]]
  [[ 0. 18. 36. 54.]   [ 0. 19. 38. 57.]   [ 0. 20. 40. 60.]]
  [[ 0. 21. 42. 63.]   [ 0. 22. 44. 66.]   [ 0. 23. 46. 69.]]]]